# -*- coding: utf-8 -*-
import os
import json
import pandas as pd
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
from sklearn.metrics import classification_report
from tenacity import retry, stop_after_attempt, wait_fixed
from settings import path

load_dotenv()
id2name = {i: line.strip() for i, line in enumerate(open(path.path_class_txt, encoding="utf-8"))}
name2id = {v: k for k, v in id2name.items()}
llm = ChatOpenAI(
    base_url="https://gateway.ai.cloudflare.com/v1/d2cbfe461e343906da9615cbceab35c6/itcast-tmf/deepseek",
    api_key="sk-9743645c05de40389822709b1bd113e3",
    model="deepseek-chat",
    model_kwargs={"response_format": {'type': 'json_object'}}

)


@retry(stop=stop_after_attempt(3), wait=wait_fixed(2))
def invoke_llm(prompt):

    return llm.invoke(prompt)


def read_data(file_path):
    "读取tab分隔的新闻标题和标签数据"
    df = pd.read_csv(file_path, sep="\t", header=None, names=['title', 'label'], encoding='utf-8')

    return df['title'].tolist(), df['label'].tolist()


def classify_news(title: str) -> dict:
    prompt = [
        {
            "role": "system",
            "content": """
    你是一名新闻分类审核员，任务是将新闻标题分类到以下类别之一：
    finance, realty, stocks, education, science, society, politics, sports, game, entertainment。
    请根据标题内容和以下关键词与示例，匹配最相关的类别。如果标题涉及教育机构但核心是社会贡献，优先归为 society。
    返回 JSON 格式：{"category": "类别", "reason": "分类原因"}

    类别关键词与示例：
    - finance: 银行、信用卡、贷款、利率 (例: "各银行信用卡挂失费迥异")
    - realty: 房产、地价、楼盘 (例: "东5环海棠公社230平准现房")
    - stocks: 股市、股指、期货 (例: "金证顾问：过山车行情意味着什么")
    - education: 学校、考试、招生 (例: "中华女子学院仅1专业招男生")
    - science: 技术、网站、宇航 (例: "“手机钱包”亮相科博会")
    - society: 社会事件、犯罪、公益 (例: "82岁老太为学生做饭扫地44年")
    - politics: 政策、国际关系 (例: "查韦斯称愿为俄罗斯提供空军基地")
    - sports: 比赛、运动员、奥运 (例: "卡佩罗：德国脚生猛的原因")
    - game: 电子游戏、网游、电竞 (例: "《赤壁OL》攻城战硝烟又起")
    - entertainment: 明星、影视、综艺 (例: "冯德伦徐若瑄隔空传情")    """},
        {
            "role": "user",
            "content": f"新闻标题：'{title}'，请分类并说明原因。"
        }
    ]

    # 调用 LLM 获取分类结果
    response = invoke_llm(prompt)
    print(response)

    result = json.loads(response.content)
    print(result)
    # 返回分类结果，包含类别和原因
    return {
        "category": result.get("category", "society"),  # 默认 society
        "reason": result.get("reason", "未明确分类，归为社会类别")
    }


if __name__ == "__main__":
    file_path = path.path_dev2_txt

    titles, true_lables = read_data(file_path)

    pred_labels = []
    results = []
    for title in titles:
        print(f'处理标题：{title}')
        result = classify_news(title)
        category = result["category"]
        pred_label = name2id.get(category, 5)  # 默认 society 为 5
        pred_labels.append(pred_label)
        results.append({
            "title": title,
            "category": category,
            "reason": result["reason"]
        })

    # 打印分类结果
    print("\n分类结果：")
    for result in results:
        print(json.dumps(result, ensure_ascii=False, indent=2))

    # 评估性能
    print("\n评估指标：")
    report = classification_report(
        true_lables,
        pred_labels,
    )
    print(report)
